// Created by John Åkerblom 2014-11-22

#include "../h3mlib.h"
#include "parse_oa.h"
#include "parse_oa_meta_type.h"
#include "../utils/memmem.h"
#include "../utils/safe_read.h"

#include "../meta/meta_object.h"

#include "../gen/def_bodies_array.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

#ifdef _WIN32
#include <windows.h>            // DEBUG
#endif

#ifdef _MSC_VER
#pragma warning(disable:4996)   // M$ standard lib unsafe
#endif

// hash function generated by mph residing in src\gen\defs.c
extern int def_bodies_hash(const char *);

static int _add_to_oa_hash(struct H3MLIB_CTX *ctx, const char *def,
    int oa_index)
{
    struct META_OA_HASH_ENTRY *oa_hash_entry = NULL;
    int oa_body_index = 0;

    // It's ofc bad to do this downcasing all the time
    // TODO stop downcasing 24/7
    char def_lower[16] = { 0 };
    char *p = def_lower;

    strncpy(def_lower, def, sizeof(def_lower) - 1);
    for (; *p; ++p)
        *p = (char)tolower((int)*p);
    def = def_lower;

    if (0 > (oa_body_index = def_bodies_hash(def))) {
        return 1;
    } else if (oa_body_index >= sizeof(def_bodies_array) / sizeof(void *)) {
        return 1;
    } else if (0 != strcmp(def_bodies_array[oa_body_index], def)) {
        return 1;
    }

    oa_hash_entry = calloc(1, sizeof(*oa_hash_entry));
    oa_hash_entry->def = strdup(def);
    oa_hash_entry->oa_body_index = oa_body_index;
    oa_hash_entry->oa_index = oa_index;
    HASH_ADD_INT(ctx->meta.oa_hash, oa_body_index, oa_hash_entry);

    return 0;
}

int parse_oa(struct H3MLIB_CTX *ctx)
{
    struct H3M *p = &ctx->h3m;
    struct H3MLIB_META *meta = &ctx->meta;
    struct H3MLIB_PARSING *parsing = &ctx->parsing;
    int orig_off = parsing->offset;
    int saved_off = 0;
    int ret = 0;
    size_t n = 0;
    unsigned int i = 0;
    struct H3M_OA *oa = &p->oa;
    struct H3M_OA_ENTRY *oa_entry = NULL;
    struct META_OA_ENTRY *meta_entry = NULL;
    int is_custom = 0;

    SAFE_READ_SIZEOF(&oa->count, parsing)

    oa->entries = calloc(oa->count, sizeof(*(oa->entries)));
    meta->oa_entries = calloc(oa->count, sizeof(*(meta->oa_entries)));
    meta->oa_entries_count = oa->count;

    for (i = 0; i < oa->count; ++i) {
        oa_entry = &oa->entries[i];
        meta_entry = &meta->oa_entries[i];

        saved_off = parsing->offset;

        SAFE_READ_SIZEOF(&oa_entry->header.def_size, parsing)
        SAFE_ALLOC_N(oa_entry->header.def, oa_entry->header.def_size + 1,
            4096) // stack overflow in Heroes3.exe if above ~16
        SAFE_READ_N(oa_entry->header.def, oa_entry->header.def_size,
            parsing)
        SAFE_READ_SIZEOF(&oa_entry->body, parsing)

        meta_entry->size = parsing->offset - saved_off;
        // Note that even the two ever-present OAs need to be set here.
        // AVLholg0.def is used in for instance Knight of Darkness.h3m (hole at 28,65)
        parse_oa_meta_type(oa_entry, &meta_entry->type, &is_custom,
            ctx->callbacks.cb_def, ctx->callbacks.cb_data);

        if (0xFF != meta_entry->type && 0 == is_custom) {
            _add_to_oa_hash(ctx, (const char *)oa_entry->header.def, i);
        }
    }

    // Check for embedded DLL by inspecting the last OA entry
    if (NULL != oa_entry && oa_entry->header.def_size > 16
        && (NULL != memmem(oa_entry->body.unknown, sizeof(oa_entry->body.unknown),
            (const unsigned char *)"potmdehex", sizeof("potmdehex") - 1))) {
        struct shellcode_oa_jmp_to_dll_load_t *shellcode_oa = NULL;
        struct shellcode_eof_load_dll_t *shellcode_eof = NULL;
        uint64_t offset = 0;
            
        shellcode_oa = (struct shellcode_oa_jmp_to_dll_load_t *)&oa_entry->header.def[7];
        
        offset = shellcode_oa->shellcode_eof_offset + sizeof(*shellcode_eof);
        if (offset > parsing->raw_size)
            goto end;

        shellcode_eof = (struct shellcode_eof_load_dll_t *)&parsing->raw[shellcode_oa->shellcode_eof_offset];

        offset = shellcode_oa->shellcode_eof_offset + shellcode_eof->dll_offset + shellcode_eof->dll_size;
        if (offset > parsing->raw_size)
            goto end;

        // Determine target with a hack looking at the address the shellcode returns to. TODO improve
        ctx->h3m_code.target = (0x005045C1 == shellcode_eof->orig_retn)? H3M_MODEMBED_TARGET_COMPLETE
            : H3M_MODEMBED_TARGET_HDMOD;

        // DLL
        ctx->h3m_code.dll_size = shellcode_eof->dll_size;
        ctx->h3m_code.dll = malloc(ctx->h3m_code.dll_size);
        memcpy(ctx->h3m_code.dll,
            &parsing->raw[shellcode_oa->shellcode_eof_offset + shellcode_eof->dll_offset],
            shellcode_eof->dll_size);

        // Shellcode
        ctx->h3m_code.shellcode_oa = malloc(sizeof(*shellcode_oa));
        memcpy(ctx->h3m_code.shellcode_oa,
            shellcode_oa,
            sizeof(*shellcode_oa));

        ctx->h3m_code.shellcode_oa_offset = 0; // Set by h3m_modembed_write_oa_eof_jmp

        // Sign OA and exploit OA
        ctx->h3m_code.extra_oa = 2; 
        ctx->h3m.oa.count -= 2;
    }

end:
    if (NULL != ctx->callbacks.cb_parse) {
        n = parsing->offset - orig_off;
        ret = ctx->callbacks.cb_parse(orig_off, "oa", &parsing->raw[orig_off],
            n, ctx->callbacks.cb_data);
    }

    return ret;
}
